import csv
from fileIO import FileIO
from log import Log
from backtest_system import BackTest
from MLtrader import MLTrader

if __name__ == '__main__':
    io = FileIO()
    data = io.read_file('data.pkl')
    trade_days = list(data.keys())
    #加载一下上证指数、国债指数作为基准
    benchmark_info = {}
    with open('benchmark_data.csv', 'r') as f:
        reader = csv.reader(f)
        next(reader)
        for row in reader:
            benchmark_info[row[0]] = [float(row[1]), float(row[2])]
    #加载一下股票信息
    stocks_info = {}
    for day in data:
        df = data[day]
        for index, row in df.iterrows():
            if index not in stocks_info:
                stocks_info[index] = {}
            stocks_info[index][day] = row
            
    #设置参数并开始回测
    initial_money = 10000000
    tax_ratio = 0.001
    max_stock_num = 10
    adjust_period = 7
    start_time = '2015-01-01'
    end_time = '2019-12-31'


    trader = MLTrader(
        initial_money=initial_money, 
        tax_ratio=tax_ratio, 
        max_stock_num=max_stock_num,
        stocks_info=stocks_info
    )
    bt = BackTest(
        initial_money=initial_money,
        stocks_info=stocks_info,
        benchmark_info=benchmark_info,
        adjust_period=adjust_period,
        start_time=start_time,
        end_time=end_time,
        trader = trader
    )
    bt.run()
